#!/usr/bin/env python
"""
Command line utility to calculate a position specific scoring matrix
for any set of sequences.

Writes tab-delimited table with residue letters in the header line
and in each column the probability of observing each of those residues
calculated from the sequence file.

Example:-
 $ python bin/ACGTCounts.py [insequencefile].fas > PSM.txt
 or
 $ python ACGTCounts.py [insequencefile].fas -out PSM.txt


Description:
Calculates position specific scoring matrix from as sequence file.

"""
# Copyright (c) Alex Leach. 2011
# Contact: Alex Leach (albl500@york.ac.uk)
# University of York, Department of Biology,
# Wentworth Way, York YO10 4DU, United Kingdom
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 3
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have receive a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>

import sys
from Bio import SeqIO
from ssummo.cmd_options import Options

def pos_counts(inFile, fileFormat='fasta'):
    countDict = {}
    iterator = SeqIO.parse(inFile, format=fileFormat)
    resSeen = set()
    nSeqs = 0
    for seqRec in iterator:
        pos = 0
        nSeqs += 1
        for res in seqRec.seq.tostring():
            res = res.upper()
            if pos not in countDict.keys():
                countDict.update( { pos : { res : 1 } } )
            elif res in countDict[pos].keys():
                countDict[pos][res] += 1
            else:
                countDict[pos].update( { res : 1 } )
            resSeen.add(res)
            pos += 1
    resSeen = sorted( list( resSeen ) )
    return countDict, resSeen, nSeqs

def tabularise(PSMDict, nseqs, out_file, residues_seen=[]):
    if len(residues_seen) == 0:
        # residues all unique, and in order
        residues = sorted(list(set(PSMDict.values().keys())))
    else:
        residues = residues_seen
    if '-' in residues:
        del(residues[residues.index('-')])
    sys.stdout.write('\t'.join(residues) + '\n')
    for i in xrange(len(PSMDict.keys())):
        if '-' in PSMDict[i].keys():
            blankCount = PSMDict[i]['-']
        else:
            blankCount = 0
        line = []
        for residue in residues:
            if residue in PSMDict[i]:
                line.append(str( float(PSMDict[i][residue]) /
                                 float(nseqs - blankCount)) )
            else:
                line.append('0')
        out_file.write('\t'.join(line) + '\n')

def parse_args(args):
    options = Options()
    options.options = { '-format' : 'fasta' ,
                        '-in' : sys.stdin,
                        '-out' : sys.stdout}
    options.help_text = { '-format' : 'Input sequence format. [fasta]',
                          '-in'     : 'Input sequence file. [/dev/stdin]',
                          '-out'    : 'Output file for PSM. [/dev/stdout]' }
    options.multiargs = []
    options.singleargs = ['-in', '-format', '-out']
    options.parse_args(args)
    return options

if __name__ == '__main__':
    args = sys.argv[1:]
    options = parse_args( args )
    if options['-in'] != sys.stdin:
        options[ '-in' ] = file(options['-in'], 'r')
    if options['-out'] != sys.stdout:
        options[ '-out' ] = file(options['-out'], 'w')
    countDict, seen, nSeqs = pos_counts(options['-in'],
                                        fileFormat=options['-format'])
    tabularise(countDict, nSeqs, options['-out'], residues_seen=seen)
    options['-in'].close()
    options['-out'].close()
